package org.example.lanchain.store;

import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.example.lanchain.bean.QwMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

/**
 * MongoDB 聊天记忆存储实现类
 * 用于将聊天记录持久化到 MongoDB 数据库中
 */
@Component  // 标记为 Spring 组件，使其能被自动扫描并管理
public class MongoChatMemoryStore implements ChatMemoryStore {  // 实现 ChatMemoryStore 接口

    /**
     * 注入 MongoDB 操作模板
     * 用于执行 MongoDB 的增删改查操作
     */
    @Autowired
    private MongoTemplate mongoTemplate;

    /**
     * 获取指定会话 ID 的所有聊天消息
     * 
     * @param memoryId 会话的唯一标识符
     * @return 返回聊天消息列表，如果没有找到则返回空列表
     */
    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        // 创建查询条件：根据 memoryId 字段进行匹配
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        // 创建查询对象
        Query query = new Query(criteria);
        // 执行查询，返回单个结果
        QwMessage qwMessage = mongoTemplate.findOne(query, QwMessage.class);
        // 如果没有找到记录，返回空列表
        if (qwMessage == null) {
            return new LinkedList<>();
        }

        // 将 JSON 格式的消息内容反序列化为消息对象列表
        return ChatMessageDeserializer.messagesFromJson(qwMessage.getContent());
    }

    /**
     * 更新指定会话 ID 的聊天消息
     * 
     * @param memoryId 会话的唯一标识符
     * @param messages 要更新的聊天消息列表
     */
    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        // 创建查询条件：根据 memoryId 字段进行匹配
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        // 创建查询对象
        Query query = new Query(criteria);

        // 过滤和清理消息内容
        List<ChatMessage> cleanedMessages = new ArrayList<>();
        
        for (ChatMessage message : messages) {
            ChatMessage cleanedMessage = cleanMessage(message);
            if (cleanedMessage != null && 
                (!(cleanedMessage instanceof AiMessage) || 
                 !((AiMessage) cleanedMessage).text().trim().isEmpty())) {
                cleanedMessages.add(cleanedMessage);
            }
        }

        // 创建更新操作对象
        Update update = new Update();
        // 设置要更新的字段：将清理后的消息列表序列化为 JSON 后存储
        update.set("content", ChatMessageSerializer.messagesToJson(cleanedMessages));
        // 执行更新操作：如果记录存在则更新，不存在则插入
        mongoTemplate.upsert(query, update, QwMessage.class);
    }

    /**
     * 清理消息内容，移除RAG检索增强的信息
     * 
     * @param message 原始消息
     * @return 清理后的消息
     */
    private ChatMessage cleanMessage(ChatMessage message) {
        if (message instanceof UserMessage) {
            UserMessage userMessage = (UserMessage) message;
            String content = userMessage.singleText();
            
            // 检查是否包含RAG增强信息
            if (content.contains("\n\nAnswer using the following information:")) {
                // 提取原始用户消息（RAG增强信息之前的部分）
                String originalContent = content.split("\n\nAnswer using the following information:")[0].trim();
                return UserMessage.from(originalContent);
            }
            
            return message;
        } else if (message instanceof AiMessage) {
            // AI消息直接返回，不需要清理
            return message;
        }
        
        // 其他类型消息也直接返回
        return message;
    }

    /**
     * 删除指定会话 ID 的所有聊天消息
     * 
     * @param memoryId 会话的唯一标识符
     */
    @Override
    public void deleteMessages(Object memoryId) {
        // 创建查询条件：根据 memoryId 字段进行匹配
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        // 创建查询对象
        Query query = new Query(criteria);
        // 执行删除操作
        mongoTemplate.remove(query, QwMessage.class);
    }
}
